from typing import List, Dict, Any, Optional
import numpy as np
import os
from kg_retrieval.env_config import MODEL_CONFIG

class OptimizedReranker:
    def __init__(self, model_path=None, use_fp16=None, model_type=None):
        # 使用环境变量或参数
        self.model_type = model_type or MODEL_CONFIG["RERANK_MODEL_TYPE"]
        self.offline_mode = MODEL_CONFIG["OFFLINE_MODE"]
        
        try:
            # 在离线模式下，跳过模型加载
            if self.offline_mode:
                print(f"运行在离线模式，使用简单算法替代重排序")
                self.reranker = None
                return
            
            if self.model_type == "flag":
                # 使用FlagEmbedding的重排序模型
                self.model_path = model_path or MODEL_CONFIG["RERANK_MODEL_PATH_FLAG"]
                self.use_fp16 = use_fp16 if use_fp16 is not None else MODEL_CONFIG["RERANK_MODEL_USE_FP16"]
                print(f"加载Flag重排序模型: {self.model_path}")
                
                try:
                    from FlagEmbedding import FlagReranker
                    # 判断是否为本地路径
                    if os.path.exists(self.model_path):
                        # 本地模型路径
                        self.reranker = FlagReranker(self.model_path, use_fp16=self.use_fp16, trust_remote_code=True)
                    else:
                        # HuggingFace模型
                        self.reranker = FlagReranker(self.model_path, use_fp16=self.use_fp16)
                except ImportError:
                    print("警告：FlagEmbedding未安装，将尝试使用SentenceTransformer")
                    self.model_type = "sentence_transformer"
            
            if self.model_type == "sentence_transformer":
                # 使用SentenceTransformer重排序模型(如Qwen3-Reranker)
                self.model_path = model_path or MODEL_CONFIG["RERANK_MODEL_PATH_ST"]
                self.use_fp16 = use_fp16 if use_fp16 is not None else MODEL_CONFIG["ST_RERANK_MODEL_USE_FP16"]
                print(f"加载SentenceTransformer重排序模型: {self.model_path}")
                
                try:
                    from sentence_transformers import CrossEncoder
                    # 加载CrossEncoder模型
                    self.reranker = CrossEncoder(self.model_path)
                    
                    # 设置为FP16模式（如果配置了的话）
                    if self.use_fp16:
                        try:
                            self.reranker.model.half()
                            print("已启用FP16模式")
                        except:
                            print("无法启用FP16模式，继续使用默认精度")
                except ImportError:
                    print("警告：sentence_transformers未安装，将使用简单算法替代重排序")
                    self.reranker = None
        except Exception as e:
            print(f"重排序模型加载失败: {e}，将使用简单算法替代")
            self.reranker = None
    
    @classmethod
    def preload(cls):
        """预加载模型，返回实例。用于启动时提前初始化模型"""
        instance = cls()
        # 预热模型，使用一个简单的查询进行测试，确保模型已经完全加载
        if instance.reranker:
            try:
                test_query = "测试查询"
                test_doc = "这是一个测试文档"
                
                if instance.model_type == "flag":
                    # 预热Flag模型
                    instance.reranker.compute_score([[test_query, test_doc]])
                elif instance.model_type == "sentence_transformer":
                    # 预热SentenceTransformer模型
                    instance.reranker.predict([[test_query, test_doc]])
                
                print(f"重排序模型预热完成: {instance.model_path}")
            except Exception as e:
                print(f"重排序模型预热失败: {e}")
        
        return instance
    
    def rerank(self, query: str, results: List[Dict[str, Any]], top_n: int = 10) -> List[Dict[str, Any]]:
        """优化的重排序函数"""
        # 只处理前 top_n 个结果
        results_to_rerank = results[:min(len(results), top_n)]
        
        if not results_to_rerank:
            return []
        
        if self.reranker:
            try:
                if self.model_type == "flag":
                    # 使用FlagReranker进行重排序
                    # 构建重排序对
                    pairs = []
                    for result in results_to_rerank:
                        # 只使用最重要的两个特征：标题和摘要
                        pairs.append([query, result["_source"]["title"]])
                        pairs.append([query, result["_source"]["summary_text"]])
                    
                    # 批量计算分数
                    scores = self.reranker.compute_score(pairs, normalize=True)
                    
                    # 为每个结果分配重排序分数 (平均标题和摘要分数)
                    for i, result in enumerate(results_to_rerank):
                        title_score = scores[i*2]
                        summary_score = scores[i*2+1]
                        # 加权平均
                        result["rerank_score"] = (title_score * 0.6) + (summary_score * 0.4)
                
                elif self.model_type == "sentence_transformer":
                    # 使用SentenceTransformer CrossEncoder进行重排序
                    pairs = []
                    for result in results_to_rerank:
                        # 将查询和文档内容组合为对
                        title = result["_source"].get("title", "")
                        summary = result["_source"].get("summary_text", "")
                        # 组合标题和摘要
                        doc_text = f"{title} {summary}"
                        pairs.append([query, doc_text])
                    
                    # 批量计算分数
                    scores = self.reranker.predict(pairs)
                    
                    # 为每个结果分配重排序分数
                    for i, result in enumerate(results_to_rerank):
                        result["rerank_score"] = float(scores[i])
            except Exception as e:
                print(f"重排序失败: {e}")
                # 如果重排序失败，回退到简单的相似度分数
                self._fallback_scoring(query, results_to_rerank)
        else:
            # 使用简单的回退方法进行评分
            self._fallback_scoring(query, results_to_rerank)
        
        # 按重排序分数排序
        reranked_results = sorted(results_to_rerank, key=lambda x: x.get("rerank_score", 0.0), reverse=True)
        
        return reranked_results
    
    def _fallback_scoring(self, query: str, results: List[Dict[str, Any]]):
        """当无法使用重排序模型时的回退评分方法"""
        query_terms = set(query.lower().split())
        
        for result in results:
            title = result["_source"].get("title", "").lower()
            summary = result["_source"].get("summary_text", "").lower()
            
            # 简单的词匹配分数
            title_matches = sum(1 for term in query_terms if term in title)
            summary_matches = sum(1 for term in query_terms if term in summary)
            
            # 标题匹配给予更高权重
            score = (title_matches * 0.6) / max(1, len(query_terms)) + (summary_matches * 0.4) / max(1, len(query_terms))
            
            # 将分数标准化到0-1范围
            result["rerank_score"] = min(max(score, 0.0), 1.0)
            
            # 如果ES提供了原始分数，也考虑进来
            if "_score" in result:
                result["rerank_score"] = result["rerank_score"] * 0.7 + min(result["_score"] / 10.0, 0.3) 